package com.ruoyi.wlxz.websocket.handler;

import com.alibaba.dashscope.app.Application;
import com.alibaba.dashscope.app.ApplicationParam;
import com.alibaba.dashscope.app.ApplicationResult;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.util.UriComponentsBuilder;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

public class MyWebSocketHandler extends TextWebSocketHandler {

    // 存储所有活跃的WebSocket会话（线程安全）
    private static final List<WebSocketSession> sessions = new CopyOnWriteArrayList<>();

    // 存储“用户ID-会话”的映射（线程安全）
    private static final Map<String, WebSocketSession> userSessionMap = new ConcurrentHashMap<>();
    // 可选：存储所有会话（用于广播）
    private static final CopyOnWriteArrayList<WebSocketSession> allSessions = new CopyOnWriteArrayList<>();

    // 连接建立时触发
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String path = session.getUri().getPath(); // 路径为 /ws/123
        String userId = UriComponentsBuilder.fromPath(path).build().getPathSegments().get(1); // 获取第2段（索引1）

        // 2. 存储用户与会话的映射
        userSessionMap.put(userId, session);
        allSessions.add(session);

        System.out.println("用户 " + userId + " 建立连接，会话ID：" + session.getId());
        System.out.println("当前在线用户数：" + userSessionMap.size());
        session.sendMessage(new TextMessage("用户 " + userId + " 连接成功！"));
    }

    // 接收客户端消息时触发
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        System.out.println("收到消息：" + payload);

        // 仅向发送消息的客户端返回响应（不广播）
        // session就是发送者的会话对象，直接调用sendMessage即可
        session.sendMessage(new TextMessage("地瓜地瓜,我是土豆,我已收到你的消息：" + payload));
        ask(String.valueOf(message), session);


    }

    // 连接关闭时触发
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        sessions.remove(session);
        System.out.println("连接关闭，当前连接数：" + sessions.size());
    }

    // 发生错误时触发
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.err.println("WebSocket错误：" + exception.getMessage());
    }

    public void ask(String inputValue, WebSocketSession session) throws IOException, NoApiKeyException, InputRequiredException {
        ResponseBodyEmitter emitter = new ResponseBodyEmitter(-1L); // 无限超时
        // 设置响应头
        emitter.send("\n"); // 发送一个初始换行符以启动连接
        ApplicationParam param = ApplicationParam.builder()
                .apiKey("sk-dd83a9bc21f5440491fe836b7aa33cee")//主机的apikey
                .appId("36e4407638164b34b63cf588d956835f")//智能体应用id,后续需要调教一个优秀的智能体
                .prompt(inputValue)//prompt提示词
                .memoryId("d2319ab3076a4beca13f96349449e980")//用户特定id,通过用户登录账号,在数据库中找到唯一的memoryid
                .incrementalOutput(true)//是否使用流式输出,我记得是这个功能来着,有些忘记了
                .build();
        Application application = new Application();
        Flowable<ApplicationResult> result = application.streamCall(param);
        result.subscribe(
                data -> {
                    try {
                        String text = data.getOutput().getText();
                        session.sendMessage(new TextMessage("我现在给你回复：" + text));
                        System.out.println("数据为" + text);
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                },
                error -> {
                    emitter.completeWithError(error);
                },
                () -> {
                    try {
                        emitter.send("[DONE]");
                        emitter.complete();
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                }
        );
    }
}